// Copyright (c) 2011 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
// Author: Markus Gutschke

#include <errno.h>
#include <poll.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <unistd.h>

#include "check.h"
#include "connection.h"

#define NOCHECK(x) do { if (x) { } } while (0)
#define min(a,b) ({ typeof(a) _a=(a); typeof(b) _b=(b); _a < _b ? _a : _b; })


struct Connection *newConnection(int fd) {
  struct Connection *conn;
  check(conn = malloc(sizeof(struct Connection)));
  initConnection(conn, fd);
  return conn;
}

void initConnection(struct Connection *conn, int fd) {
  memset(conn, 0, sizeof(struct Connection));
  conn->fd = fd;
}

void destroyConnection(struct Connection *conn) {
  if (conn) {
    if (conn->fd >= 0) {
      if (conn->ssl) {
        if (conn->handle) {
          closeSSLHandle(conn->handle);
        }
        deleteSSL(conn->ssl);
      } else {
        check(!conn->handle);
      }
      check(!close(conn->fd));
    } else {
      check(!conn->handle);
      check(!conn->ssl);
    }
    conn->fd = -1;
    free(conn->buf);
  }
}

void deleteConnection(struct Connection *conn) {
  destroyConnection(conn);
  free(conn);
}

void upgradeConnection(struct Connection *conn) {
  check(conn->fd >= 0);
  check(!conn->ssl);
  check(!conn->handle);
  conn->ssl = newSSL();
  conn->handle = openSSLHandle(conn->ssl, conn->fd);
}

int connectionIsSecure(struct Connection *conn) {
  return conn && conn->handle;
}

#ifdef VERBOSE
static void writeCtrl(int fd, const char *buf, size_t sz) {
  char *cpy;
  check(cpy = malloc(sz));
  for (size_t i = 0; i < sz; ++i) {
    cpy[i] = (buf[i] >= ' ' && buf[i] < 0x7F) || buf[i] == '\n' ?
              buf[i] : '.';
  }
  NOCHECK(write(fd, cpy, sz));
  free(cpy);
}
#endif

int connIsReadable(struct Connection *conn, int timeout) {
  if (conn->buf) {
    return 1;
  }
  if (conn->handle) {
    check(conn->buf = malloc(conn->sz = 1));
    conn->offset = 0;
    ssize_t rc = SSLread(conn->handle, conn->buf, conn->sz);
    if (rc == 1) {
#ifdef VERBOSE
      NOCHECK(write(2, conn->handle ? "\033[43m" : "\033[42m", 5));
      writeCtrl(2, conn->buf, rc);
      NOCHECK(write(2, "\033[m\n", 4));
#endif
      return 1;
    }
    free(conn->buf);
    conn->buf = NULL;
    conn->sz = 0;
  }
  struct pollfd fds = {
    .fd     = conn->fd,
    .events = POLLIN,
  };
  return TEMP_FAILURE_RETRY(poll(&fds, 1, timeout)) == 1 &&
         (fds.revents & POLLIN);
}

int connIsWriteable(struct Connection *conn, int timeout) {
  struct pollfd fds = {
    .fd     = conn->fd,
    .events = POLLOUT,
  };
  return TEMP_FAILURE_RETRY(poll(&fds, 1, timeout)) == 1 &&
         (fds.revents & POLLOUT);
}

ssize_t connRead(struct Connection *conn, char *buf, size_t sz) {
  return connReadTMO(conn, buf, sz, -1);
}

ssize_t connReadTMO(struct Connection *conn, char *buf, size_t sz, int tmo) {
  ssize_t rc;
  for (;;) {
    if (conn->buf) {
      rc = min(sz, conn->sz);
      memcpy(buf, conn->buf + conn->offset, rc);
      if (!(conn->sz -= rc)) {
        free(conn->buf);
        conn->buf = NULL;
        ssize_t xtra;
        if (conn->handle) {
          xtra = SSLread(conn->handle, buf + rc, sz - rc);
        } else {
          xtra = read(conn->fd, buf + rc, sz - rc);
        }
        if (xtra > 0) {
#ifdef VERBOSE
          NOCHECK(write(2, conn->handle ? "\033[43m" : "\033[42m", 5));
          writeCtrl(2, buf + rc, xtra);
          NOCHECK(write(2, "\033[m\n", 4));
#endif
          rc += xtra;
        }
      } else {
        conn->offset += rc;
      }
      return rc;
    }
    if (conn->handle) {
      rc = SSLread(conn->handle, buf, sz);
    } else {
      rc = read(conn->fd, buf, sz);
    }
    if (rc < 0) {
      if (errno == EINTR) {
        continue;
      } else if (errno == EAGAIN) {
        if (!tmo || !connIsReadable(conn, tmo)) {
          errno = EAGAIN;
          return -1;
        }
        continue;
      }
    }
#ifdef VERBOSE
    if (rc > 0) {
      NOCHECK(write(2, conn->handle ? "\033[43m" : "\033[42m", 5));
      writeCtrl(2, buf, rc);
      NOCHECK(write(2, "\033[m\n", 4));
    }
#endif
    return rc;
  }
}

ssize_t connReadAll(struct Connection *conn, char *buf, size_t sz) {
  return connReadAllTMO(conn, buf, sz, -1);
}

ssize_t connReadAllTMO(struct Connection *conn, char *buf, size_t sz, int tmo){
  ssize_t rc = 0;
  while (sz > 0) {
    ssize_t nbytes = connReadTMO(conn, buf, sz, tmo);
    if (nbytes <= 0) {
      if (rc) {
        break;
      }
      return nbytes;
    }
    rc  += nbytes;
    buf += nbytes;
    sz  -= nbytes;
  }
  return rc;
}

ssize_t connPeek(struct Connection *conn, char *buf, size_t sz) {
  ssize_t rc;
  if (conn->buf) {
    rc = min(sz, conn->sz);
    memcpy(buf, conn->buf + conn->offset, rc);
    return rc;
  }
  for (;;) {
    if (conn->handle) {
      // In SSL mode, we don't generally have the ability to peek into the
      // stream. Emulate peeking by extra buffering.
      rc = connRead(conn, buf, sz);
      if (rc > 0) {
        check(conn->buf = malloc(rc));
        memcpy(conn->buf, buf, rc);
        conn->sz = rc;
        conn->offset = 0;
      }
      return rc;
    }
    rc = recv(conn->fd, buf, sz, MSG_PEEK);
    if (rc < 0) {
      if (errno == EINTR) {
        continue;
      } else if (errno == EAGAIN) {
        connIsReadable(conn, -1);
        continue;
      }
    }
    return rc;
  }
}

ssize_t connWrite(struct Connection *conn, const char *buf, size_t sz) {
  return connWriteTMO(conn, buf, sz, -1);
}

ssize_t connWriteTMO(struct Connection *conn, const char *buf, size_t sz,
                     int tmo) {
  for (size_t ret = 0;;) {
    ssize_t rc;
    if (conn->handle) {
      rc = SSLwrite(conn->handle, buf, sz);
    } else {
      rc = write(conn->fd, buf, sz);
    }
    if (rc < 0) {
      if (errno == EINTR) {
        continue;
      } else if (errno == EAGAIN) {
        if (!tmo || !connIsWriteable(conn, tmo)) {
          errno = EAGAIN;
          if (ret) {
            return ret;
          }
          return -1;
        }
        continue;
      }
    } else if (rc > 0) {
#ifdef VERBOSE
    NOCHECK(write(2, conn->handle ? "\033[44m" : "\033[45m", 5));
    writeCtrl(2, buf, rc);
    NOCHECK(write(2, "\033[m\n", 4));
#endif
      buf += rc;
      sz  -= rc;
      ret += rc;
      if (sz > 0) {
        continue;
      }
    }
    return ret;
  }
}

ssize_t connWriteV(struct Connection *conn, struct iovec *io, int cnt) {
  return connWriteVTMO(conn, io, cnt, -1);
}

ssize_t connWriteVTMO(struct Connection *conn, struct iovec *io, int cnt,
                      int tmo) {
  ssize_t rc = 0;
  for (int i = 0; i < cnt; ++i) {
    if (!io[i].iov_len) {
      continue;
    }
    ssize_t nbytes = connWriteTMO(conn, io[i].iov_base, io[i].iov_len, tmo);
    if (nbytes <= 0) {
      if (rc) {
        break;
      }
      return nbytes;
    }
    rc += nbytes;
    if (nbytes != io[i].iov_len) {
      break;
    }
  }
  return rc;
}
